package com.gitee.luosl.cola.lang;

import java.util.*;

public class GroupedIterator<T> extends AbstractIterator<List<T>> {

    private final Iterator<T> that;
    private final int size;
    private final int step;
    private final int gap;
    private final int preserve;

    private Object[] buffer;

    public GroupedIterator(Iterator<T> that, int size, int step) {
        if (size < 1 || step < 1) {
            throw new IllegalArgumentException("size (" + size + ") and step (" + step + ") must both be positive");
        }
        this.that = that;
        this.size = size;
        this.step = step;
        this.gap = Math.max(step - size, 0);
        this.preserve = Math.max(size - step, 0);
        this.buffer = take(that, new Object[size], 0, size);
    }

    @Override
    public boolean hasNext() {
        return buffer.length > 0;
    }

    @Override
    @SuppressWarnings("unchecked")
    protected List<T> getNext() {
        if (buffer.length == 0) {
            throw new NoSuchElementException();
        }
        final Object[] result = buffer;
        if (that.hasNext()) {
            buffer = new Object[size];
            if (preserve > 0) {
                System.arraycopy(result, step, buffer, 0, preserve);
            }
            if (gap > 0) {
                drop(that, gap);
                buffer = take(that, buffer, preserve, size);
            } else {
                buffer = take(that, buffer, preserve, step);
            }
        } else {
            buffer = new Object[0];
        }

        return new ArrayList(Arrays.asList(result));
    }

    private static void drop(Iterator<?> source, int count) {
        for (int i = 0; i < count && source.hasNext(); i++) {
            source.next();
        }
    }

    private static Object[] take(Iterator<?> source, Object[] target, int offset, int count) {
        int i = offset;
        while (i < count + offset && source.hasNext()) {
            target[i] = source.next();
            i++;
        }
        if (i < target.length) {
            final Object[] result = new Object[i];
            System.arraycopy(target, 0, result, 0, i);
            return result;
        } else {
            return target;
        }
    }
}